import bpy
import os, shutil
import numpy as np
from argparse import ArgumentParser

from texture_allocator import TextureAllocator
from generate_background import draw_background

parser = ArgumentParser(description='bpy arg parser')
parser.add_argument("--cuda_id", type=int, default=0)
parser.add_argument("--mesh_type", type=str, default='mesh_mip')
parser.add_argument("--mesh_root_dir", type=str, default='mesh_data')
parser.add_argument("--cam_pose_dir", type=str, default='cam_pose')
parser.add_argument("--save_root_dir", type=str, default='render_res')
parser.add_argument("--mesh_list", nargs="+", type=str, default=[])
parser.add_argument("--is_texture", action='store_true')
parser.add_argument('--debug_mode', type=int, default=10)
parser.add_argument('--debug_video_step', type=int, default=5)
# parser.add_argument('--debug_id', type=int, default=0)
parser.add_argument('--write_cover', action='store_true')
parser.add_argument('--fov_scale', type=float, default=1.0)
args = parser.parse_args()
# light_energy = (100.0, 60.0, 60.0)

if args.mesh_type == 'mesh_mip':
    from render_cfgs.mip.light import light_cfg
    from render_cfgs.mip.background import back_sphere_radius
else:
    assert args.mesh_type == 'mesh_dtu'
    from render_cfgs.dtu.light import light_cfg
    from render_cfgs.dtu.background import back_sphere_radius

# if args.debug_mode > 0:
    # args.save_root_dir = os.path.join(args.save_root_dir, 'debug' + str(args.debug_id))
    # args.write_cover = True
if args.is_texture:
    light_cfg = light_cfg['texture']
    texture_allocator = TextureAllocator(bpy)
    save_root_dir = os.path.join(args.save_root_dir, 'texture')
else:
    light_cfg = light_cfg['mesh']
    save_root_dir = os.path.join(args.save_root_dir, 'mesh')

bpy.data.objects['Light'].location = light_cfg['pose'][0]
bpy.data.lights['Light'].energy = light_cfg['energy'][0]
bpy.ops.object.light_add(type="POINT", radius=1, location=light_cfg['pose'][1])
bpy.data.objects['Point'].data.energy = light_cfg['energy'][1]
bpy.ops.object.light_add(type="POINT", radius=1, location=light_cfg['pose'][2])
bpy.data.objects['Point.001'].data.energy = light_cfg['energy'][2]

# draw the sphere background
draw_background(args.is_texture)
bpy.data.objects['SurfSphere'].scale = (back_sphere_radius, back_sphere_radius, back_sphere_radius)
bpy.data.objects['SurfSphere'].location = (0.0, 0, 0)
bpy.data.objects['SurfSphere'].rotation_euler = (0, 0, 0)

# set the render configuration for Cycles
bpy.data.scenes['Scene'].render.engine = 'CYCLES'
bpy.data.scenes['Scene'].cycles.samples = 1024
bpy.data.scenes['Scene'].cycles.time_limit = 5
bpy.context.scene.cycles.device = 'GPU'
bpy.context.preferences.addons["cycles"].preferences.get_devices()
bpy.context.preferences.addons['cycles'].preferences.compute_device_type = 'CUDA'

for idx, d in enumerate(bpy.context.preferences.addons["cycles"].preferences.devices):
    d["use"] = 0
    if d["name"].startswith("NVIDIA") and idx == args.cuda_id:
        d["use"] = 1
        # break

bpy.data.objects.remove(bpy.data.objects['Cube'])

mesh_dir = os.path.join(args.mesh_root_dir, args.mesh_type)
if len(args.mesh_list) == 0:
    mesh_list = os.listdir(mesh_dir)
else:
    mesh_list = args.mesh_list

for mesh in mesh_list:
    # if mesh == 'room':
    #     bpy.data.objects['SurfSphere'].scale = (15, 15, 15)
    bpy.ops.wm.ply_import(filepath=os.path.join(mesh_dir, mesh, "mesh.ply"))
    if args.is_texture:
        texture_allocator.init_texture()
        texture_allocator.set_texture()

    bpy.data.objects['Camera'].rotation_mode = 'QUATERNION'
    bpy.data.cameras['Camera'].lens_unit = 'FOV'
    
    traj = 'scan24' if mesh.startswith('scan') else mesh
    mesh_name = mesh.split('.')[0] # "bicycle".split('.')[0] = "bicycle", okay
    extr_dir = os.path.join(args.cam_pose_dir, args.mesh_type, traj, 'extrinsics')
    intr_dir = os.path.join(args.cam_pose_dir, args.mesh_type, traj, 'intrinsics')
    save_dir = os.path.join(save_root_dir, args.mesh_type, mesh)

    if os.path.exists(save_dir) and not args.write_cover:
        assert False, 'The save directory already exists.'
    os.makedirs(save_dir, exist_ok=True)
    
    # render all the image for the mesh of type "mesh_type"
    img_list = sorted(os.listdir(intr_dir))
    for img_id, img in enumerate(img_list):
        if args.debug_mode > 0 and img_id % args.debug_video_step != 0:
            continue

        img_name = img.split('.')[0]
        data_rot = np.loadtxt(os.path.join(extr_dir, img_name + '_q.txt'))
        data_loc = np.loadtxt(os.path.join(extr_dir, img_name + '_t.txt'))
        data_intr = np.loadtxt(os.path.join(intr_dir, img))
        data_name = img_name + '.png'
        
        for i in range(3):
            bpy.data.objects['Camera'].rotation_quaternion[i] = data_rot[i]
            bpy.data.objects['Camera'].location[i] = data_loc[i]
        bpy.data.objects['Camera'].rotation_quaternion[3] = data_rot[3]

        # bpy.data.cameras['Camera'].lens = data_intr[0]
        bpy.data.cameras['Camera'].angle = data_intr[0] * args.fov_scale
        if data_intr[1] > 1:
            bpy.data.scenes['Scene'].render.pixel_aspect_x = data_intr[1]
        else:
            bpy.data.scenes['Scene'].render.pixel_aspect_x = 1
            bpy.data.scenes['Scene'].render.pixel_aspect_y = 1 / data_intr[1]
        bpy.data.scenes['Scene'].render.pixel_aspect_y = 1
        downsample_rate1 = 1
        bpy.data.scenes['Scene'].render.resolution_x = int(data_intr[2]/downsample_rate1)
        bpy.data.scenes['Scene'].render.resolution_y = int(data_intr[3]/downsample_rate1)

        bpy.ops.render.render() # render an image
        bpy.data.images["Render Result"].save_render(filepath=os.path.join(save_dir, data_name)) # save the image

        if args.debug_mode > 0 and img_id > args.debug_mode * args.debug_video_step:
            break
    
    # clean up
    # bpy.data.objects.remove(bpy.data.objects[mesh_name])
    bpy.data.objects.remove(bpy.data.objects["mesh"])
    bpy.ops.outliner.orphans_purge(do_recursive=True)